package cn.kennylee.codehub.springboot.microtrace;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.apache.dubbo.rpc.Invocation;
import org.slf4j.MDC;
import org.springframework.util.StringUtils;

import static cn.kennylee.codehub.springboot.microtrace.TraceUtils.getHeaderTraceIdName;


/**
 * 链路跟踪上下文
 * <p>Created on 2025/5/7.</p>
 *
 * @author kennylee
 * @since 0.0.1
 */
@Slf4j
public class TraceContext {

    private static final ThreadLocal<String> TRACE_ID = new ThreadLocal<>();

    /**
     * 从HTTP请求头初始化（Web入口）
     * <ol>
     *  <li>从HTTP请求头中获取Trace ID</li>
     *  <li>如果请求头中没有Trace ID，则生成一个新的Trace ID</li>
     *  <li>将Trace ID存储在ThreadLocal中</li>
     * </ol>
     *
     * @param request HTTP请求
     */
    public static void initFromHttpHeaders(HttpServletRequest request) {

        String traceId = request.getHeader(getHeaderTraceIdName());

        // 如果请求头中没有 Trace ID，尝试从查询参数
        TraceProperties traceProperties = SpringUtil.getBean(TraceProperties.class);
        if (!StringUtils.hasLength(traceId) && traceProperties.isAcceptQueryTraceId()) {
            traceId = StrUtil.trimToNull(request.getParameter(traceProperties.getMdcTraceIdName()));
        }

        // 处理是否接收传递的 Trace ID
        traceId = TraceUtils.acceptTransferTraceId(traceId);

        // 如果没有 Trace ID，生成一个新的
        if (!StringUtils.hasLength(traceId)) {
            // 生成新的 Trace ID
            traceId = generateTraceId();
            log.debug("从http请求头找不到trace id, 新生成 Trace ID: {}", traceId);
        }
        TRACE_ID.set(traceId);
    }

    /**
     * 从Dubbo隐式参数初始化（RPC入口）
     * <ol>
     *     <li>从Dubbo隐式参数中获取Trace ID</li>
     *     <li>如果有 Trace ID，则将其存储在ThreadLocal中</li>
     * </ol>
     *
     * @param invocation Dubbo调用上下文
     */
    public static void initFromRpcContext(Invocation invocation) {
        String traceId = invocation.getAttachment(getHeaderTraceIdName());

        if (StringUtils.hasLength(traceId)) {
            TRACE_ID.set(traceId);
        }
    }

    /**
     * 获取当前Trace ID（应用内使用）
     *
     * @return 当前Trace ID
     */
    public static String getTraceId() {
        return TRACE_ID.get();
    }

    /**
     * 设置Trace ID（应用内使用）
     *
     * @param traceId Trace ID
     */
    public static void setTraceId(String traceId) {
        log.debug("设置Trace ID: {}", traceId);
        TRACE_ID.set(traceId);
    }

    /**
     * 清理上下文，包括MDC和ThreadLocal
     */
    public static void clear() {
        // 请求完成后清理 MDC
        MDC.clear();
        TRACE_ID.remove();
    }

    /**
     * 生成新的Trace ID
     *
     * @return 新的Trace ID
     */
    public static String generateTraceId() {
        return IdUtil.fastSimpleUUID();
    }
}
